import sys
import random
import pylab as plt
import numpy as np
import math
import os
import networkx as nx

from params import *

datasets = ['bkite','fsquare','gowalla']
lbls = ['Brightkite', 'Foursquare', 'Gowalla']
node_dist = [5651.0,8494.0,5663.0]
link_dist = [2041.0,1442.0,1792.0]
degs = [7.88,22.07,9.48,7.0]

def ecdf(data):
    l1 = math.log10(min(data))
    l2 = math.log10(max(data)) 
    bins = 10**np.linspace(l1,l2, 100)

    pdf1,bins1,x = plt.hist(data,
            bins=bins,cumulative=False,normed=False)
    plt.close()

    c = [pdf1[0]]
    for j in range(1,len(pdf1)):
        t = c[j-1]
        c.append(t+pdf1[j])

    total = c[-1]
    c1 = map(lambda x: 1.0*x/total, c)
    print 'total ', total
    return bins1[:-1],c1

def load_graph(file):
    s = 0.0
    k = 0
    i = 0
    g = nx.Graph()
    for line in open(file):
        i += 1
        if i % 1000000 == 0:
            print 'line ', i
        if line.startswith('#'):
            continue
        u1,u2,dist = line.split(' ')
        u1,u2 = map(int,(u1,u2))
        dist = float(dist)
        dist = max((1.0,dist))
        s += dist
        k += 1
        g.add_edge(u1,u2,weight=dist)
    avg_dist = s/k
    return g, avg_dist

def user_shape_factor(file,graph):
    tr = 0
    user_shapes = {}
    for line in open(file):
        #if random.random() >= 0.01:
        #    continue
        n1,n2,n3 = sorted(map(int,line.split()))
        l1 = graph[n1][n2]['weight']
        l2 = graph[n1][n3]['weight']
        l3 = graph[n2][n3]['weight']
        l = float(l1+l2+l3)/3
        user_shapes.setdefault(n1,[]).append(l)
        user_shapes.setdefault(n2,[]).append(l)
        user_shapes.setdefault(n3,[]).append(l)

    data = {}
    for u in user_shapes:
        k = graph.degree(u)
        dist = sum(user_shapes[u])/len(user_shapes[u])
        i = math.exp(round(math.log(k),1))
        data.setdefault(i,[]).append(dist)

    x = sorted(data)
    y = [float(sum(data[k]))/len(data[k]) for k in x]
    return x,y

plots = []
for dataset in datasets:
    geo_trace_file = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,
            '%s_geo_null_graph_1.txt'%dataset)
    social_trace_file = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,
            '%s_social_null_graph_2.txt'%dataset)
    social_triangle_file = os.path.join(WORKDIR, 'gsn','results',dataset,'%s_triangles.txt'%dataset)
    geo_triangle_file = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,'%s_geo_null_1_triangles.txt'%dataset)
    print 'Dataset %s'%dataset
    trace_file = os.path.join(WORKDIR, 'gsn','traces',dataset,
            '%s_graph.txt'%dataset)
    triangle_file = os.path.join(WORKDIR, 'gsn','results',dataset,'%s_triangles.txt'%dataset)

    graph,avg_dist = load_graph(trace_file)
    print 'average link length ', avg_dist
    x,y = user_shape_factor(triangle_file,graph)
    #y = map(lambda i: i/avg_dist,y)

    del graph

    geo_graph, geo_avg_dist = load_graph(geo_trace_file)
    print "average geo link length", geo_avg_dist
    x_geo,y_geo = user_shape_factor(geo_triangle_file,geo_graph)
    #y_geo = map(lambda i: i/geo_avg_dist,y_geo)

    del geo_graph

    soc_graph, soc_avg_dist = load_graph(social_trace_file)
    print "average social link length", soc_avg_dist
    x_soc,y_soc = user_shape_factor(social_triangle_file,soc_graph)
    #y_soc = map(lambda i: i/soc_avg_dist,y_soc)

    del soc_graph
    print ''

    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    plt.loglog(x,y,'kx')
    plt.loglog(x_geo,y_geo,'k^', mfc='None')
    plt.loglog(x_soc,y_soc,'ko', mfc='None')
    plt.legend(['Original data','Geo model','Social model'],
            loc='lower right',
            numpoints=1)
    plt.ylabel(r'$\langle l_{\Delta} \rangle$')
    plt.xlabel('Degree')
    plt.grid(True)
    plt.ylim((1e1,1e4))
    plt.xlim((1e0,1e4))
    plt.savefig('%s_null_shape_degree.pdf'%dataset)
    plt.close()
